# 1导入模块
import datetime
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import RFE
from sklearn.metrics import classification_report, roc_auc_score, accuracy_score
import warnings
import joblib
import seaborn as sns

# --- 配置 ---
# 忽略一些未来版本的警告，使输出更整洁
warnings.filterwarnings('ignore', category=FutureWarning)
# 设置pandas显示选项
pd.set_option('display.max_columns', None)

# 解决中文乱码问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 2定义获取数据函数
def data_preprocess(data_path):
    # 读取数据
    df = pd.read_csv(data_path, encoding='gbk')
    # 处理缺失值
    df.fillna(0, inplace=True)
    # 删除重复行
    df.drop_duplicates(inplace=True)
    # 去除冗余数据列
    df.drop(['EmployeeNumber', '是否满18', 'StandardHours'], axis=1, inplace=True)
    # 将类别型数据转换为数值型数据
    #   对无序object类别数据进行热编码处理
    df = pd.get_dummies(df, columns=['部门', '出差频率', '性别', '工作岗位', '婚姻状况', '加班', '专业'])
    #   将热编码后的数据从布尔型转换为数值型数据
    for col in df.columns:
        if df[col].dtype == 'bool':
            df[col] = df[col].astype(int)
    #   去除冗余数据列
    df.drop(['加班_Yes', '性别_Female'], axis=1, inplace=True)
    # 返回处理后的数据
    return df


# 3数据部分可视化
def data_analysis(data):
    """
    1.查看数据整体情况
    2.查看离职总体人数
    3.查看离职情况随月薪的变化
    4.查看离职情况随工作满意度的变化
    5.查看离职情况随薪资增长百分率的变化
    :param data:数据源
    :return:
    """
    data = data.copy(deep=True)
    # 1.查看数据整体情况
    print(data.info())
    print(data.head())
    # 2.查看离职总体人数
    fig = plt.figure(figsize=(10, 30))
    ax1 = fig.add_subplot(3, 1, 1)
    #   绘制直方图以分析离职情况的分布
    ax1.hist(data['离职'], bins=4)
    ax1.set_title('离职人数直方图')
    # 3.查看离职情况随月薪的变化
    bins = [0, 4000, 8000, 12000, 16000, 20000]  # 定义工资区间
    labels = ['0-4000', '4000-8000', '8000-12000', '12000-16000', '16000-20000']  # 区间标签
    # 创建新的列 '工资区间'
    data['工资区间'] = pd.cut(data['月收入'], bins=bins, labels=labels, include_lowest=True)
    # 按工资区间统计离职人数
    grouped_data = data.groupby('工资区间')['离职'].sum().reset_index()
    # 绘制柱状图
    ax2 = fig.add_subplot(3, 1, 2)
    ax2.bar(grouped_data['工资区间'], grouped_data['离职'], color='skyblue')
    ax2.set_title('离职情况随月薪的变化')
    ax2.set_xlabel('月收入')
    ax2.set_ylabel('离职')
    # 4.查看离职情况随工作满意度的变化
    grouped_data2 = data.groupby('工作满意度')['离职'].sum().reset_index()
    ax3 = fig.add_subplot(3, 1, 3)
    ax3.bar(grouped_data2['工作满意度'], grouped_data2['离职'], color='r')
    ax3.set_title('离职情况随工作满意度的变化')
    ax3.set_xlabel('工作满意度')
    ax3.set_ylabel('离职')
    plt.savefig('../data/fig/picture.png')


# 3特征工程
def feature_engineering(data):
    """
       1.准备数据集
       2.划分训练集和测试集
       3.标准化
       :param data:
       :return:
       """
    data = data.copy(deep=True)
    X = data.drop(['离职'], axis=1)
    y = data['离职']
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    # 标准化
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)
    return X_train, X_test, y_train, y_test


# 4模型训练
def model_train(X_train, y_train):
    # 随机森林模型 - 用于特征重要性分析
    rf = RandomForestClassifier(n_estimators=200, random_state=42, class_weight='balanced')
    rf.fit(X_train, y_train)

    # 逻辑回归模型 - 用于系数解释
    lr = LogisticRegression(max_iter=1000, random_state=42, class_weight='balanced')
    lr.fit(X_train, y_train)
    return rf, lr


# 5模型评估及模型保存
def evaluate_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_proba)

    print(classification_report(y_test, y_pred))
    print(f"AUC Score: {roc_auc_score(y_test, y_proba):.4f}")
    print(f"Accuracy: {accuracy_score(y_test, y_pred):.4f}")

    # 保存模型
    joblib.dump(model, '../model/model.pkl')

    return auc




# 主函数
if __name__ == '__main__':
    # 1获取数据
    df = data_preprocess('../data/train_translate.csv')
    df.info()
    # 2数据分析
    data_analysis(df)
    # 3特征工程
    X_train, X_test, y_train, y_test = feature_engineering(df)
    # 4模型训练
    rf, lr = model_train(X_train, y_train)
    # 5模型评估
    auc = evaluate_model(lr, X_test, y_test)
    print(auc)
    
    # 展示特征相关性
    X = df.drop(['离职'], axis=1)
    y = df['离职']
    plt.figure(figsize=(12, 8))
    feature_importance = pd.Series(rf.feature_importances_, index=X.columns)
    feature_importance.sort_values(ascending=False, inplace=True)
    sns.barplot(x=feature_importance.values[:10], y=feature_importance.index[:10])
    plt.title('Top 10 Random Forest Feature Importance')
    plt.xlabel('Importance Score')
    plt.show()
